from keras.models import Model
from keras.layers.core import Dense, Dropout, Activation
from keras.layers.convolutional import Convolution2D
from keras.layers.pooling import GlobalAveragePooling2D
from keras.layers import Input
from keras.layers import BatchNormalization
from keras.regularizers import l2
import keras.backend as K
import torch
from att_blocks import *
from dense_blocks import *


class ADCNN:
    def __init__(self, img_dim, nb_classes, depth=40, nb_dense_block=3, growth_rate=12, nb_filter=16, dropout_rate=None,
                 weight_decay=1E-4, verbose=True):
        self.img_dim = img_dim
        self.nb_classes = nb_classes
        self.depth = depth
        self.nb_dense_block = nb_dense_block
        self.growth_rate = growth_rate
        self.nb_filter = nb_filter
        self.dropout_rate = dropout_rate
        self.weight_decay = weight_decay
        self.verbose = verbose
        self.build_model()

    def build_model(self):
        n_channels = 64
        model_input = Input(shape=self.img_dim)

        concat_axis = 1 if K.image_data_format() == "th" else -1

        assert (self.depth - 4) % 3 == 0, "Depth must be 3 N + 4"

        # layers in each dense block
        nb_layers = int((self.depth - 4) / 3)

        # Initial convolution
        x = Convolution2D(self.nb_filter, (3, 3), kernel_initializer="he_uniform", padding="same", name="initial_conv2D",
                          use_bias=False,
                          kernel_regularizer=l2(self.weight_decay))(model_input)

        x = BatchNormalization(axis=concat_axis, gamma_regularizer=l2(self.weight_decay),
                                beta_regularizer=l2(self.weight_decay))(x)

        # Add attention block
        x = attention_block(x, encoder_depth=1)

        # Add dense blocks
        for block_idx in range(self.nb_dense_block - 1):
            x, self.nb_filter = dense_block(x, nb_layers, self.nb_filter, self.growth_rate, dropout_rate=self.dropout_rate,
                                       weight_decay=self.weight_decay)
            # add transition_block
            x = transition_block(x, self.nb_filter, dropout_rate=self.dropout_rate, weight_decay=self.weight_decay)

        # The last dense_block does not have a transition_block
        x, self.nb_filter = dense_block(x, nb_layers, self.nb_filter, self.growth_rate, dropout_rate=self.dropout_rate,
                                   weight_decay=self.weight_decay)

        x = Activation('relu')(x)
        x = GlobalAveragePooling2D()(x)
        x = Dense(self.nb_classes, activation='softmax', kernel_regularizer=l2(self.weight_decay),
                  bias_regularizer=l2(self.weight_decay))(x)

        self.adcnn = Model(inputs=model_input, outputs=x)

        if self.verbose:
            print("ADCNN-%d-%d created." % (self.depth, self.growth_rate))

    def get_model(self):
        return self.adcnn